package com.wuwei.elearning.config.factory;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.wuwei.elearning.config.AiConfig;
import com.wuwei.elearning.question.model.ChatMessage;
import com.wuwei.elearning.question.model.ChatResponse;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.MediaType;
import org.springframework.http.client.ClientHttpResponse;
import org.springframework.web.client.ResponseExtractor;
import org.springframework.web.client.RestTemplate;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * 百炼AI客户端实现（阿里云百炼平台）
 */
@Slf4j
public class BailianClient implements AiClient {

    private final AiConfig aiConfig;
    private final RestTemplate restTemplate;
    private static final Pattern DATA_PATTERN = Pattern.compile("data: (.*)");
    private final ObjectMapper objectMapper = new ObjectMapper();
    
    public BailianClient(AiConfig aiConfig) {
        this.aiConfig = aiConfig;
        this.restTemplate = new RestTemplate();
    }
    
    @Override
    public String streamChatCompletion(List<ChatMessage> messages, SseEmitter emitter, StringBuilder responseBuilder) {
        try {
            // 将消息转换为百炼格式
            List<Map<String, Object>> bailianMessages = convertToBailianFormat(messages);
            
            // 创建请求体
            Map<String, Object> requestBody = createRequestBody(bailianMessages, true);
            
            // 设置请求头
            HttpHeaders headers = createHeaders();
            HttpEntity<Map<String, Object>> requestEntity = new HttpEntity<>(requestBody, headers);
            
            // 百炼API URL
            String url = "https://" + aiConfig.getEndpoint() + "/api/v1/services/aigc/text-generation/generation";
            log.info("发送流式请求到百炼API: {}", url);
            log.info("请求头: {}", headers);
            log.info("请求体: {}", requestBody);
            
            try {
                // 使用ResponseExtractor处理流式响应
                restTemplate.execute(url, HttpMethod.POST, 
                    req -> {
                        req.getHeaders().addAll(headers);
                        req.getHeaders().setContentType(MediaType.APPLICATION_JSON);
                        try {
                            String jsonBody = objectMapper.writeValueAsString(requestBody);
                            req.getBody().write(jsonBody.getBytes());
                        } catch (IOException e) {
                            log.error("写入请求体时出错", e);
                        }
                    },
                    new ResponseExtractor<Void>() {
                        @Override
                        public Void extractData(ClientHttpResponse response) throws IOException {
                            try (BufferedReader reader = new BufferedReader(new InputStreamReader(response.getBody()))) {
                                String line;
                                while ((line = reader.readLine()) != null) {
                                    if (line.isEmpty()) continue;
                                    
                                    if (line.startsWith("data:")) {
                                        Matcher matcher = DATA_PATTERN.matcher(line);
                                        if (matcher.find()) {
                                            String data = matcher.group(1).trim();
                                            if ("[DONE]".equals(data)) {
                                                break;
                                            }
                                            
                                            try {
                                                // 解析JSON响应 (百炼API的响应格式可能与OpenAI不同，需要根据实际情况调整)
                                                Map<String, Object> responseMap = objectMapper.readValue(data, Map.class);
                                                if (responseMap.containsKey("output") && ((Map)responseMap.get("output")).containsKey("text")) {
                                                    String content = (String) ((Map)responseMap.get("output")).get("text");
                                                    String newContent = content.substring(responseBuilder.length());
                                                    responseBuilder.append(newContent);
                                                    
                                                    ChatResponse partialResponse = new ChatResponse();
                                                    partialResponse.setContent(responseBuilder.toString());
                                                    partialResponse.setIsComplete(false);
                                                    emitter.send(partialResponse);
                                                }
                                            } catch (Exception e) {
                                                log.error("解析流式响应时出错: {}", e.getMessage(), e);
                                            }
                                        }
                                    }
                                }
                            } catch (Exception e) {
                                log.error("处理流式响应时出错", e);
                                ChatResponse errorResponse = new ChatResponse();
                                errorResponse.setContent("处理API响应时出错: " + e.getMessage());
                                errorResponse.setIsComplete(true);
                                try {
                                    emitter.send(errorResponse);
                                } catch (IOException ex) {
                                    log.error("发送错误响应时出错", ex);
                                }
                            }
                            return null;
                        }
                    });
                
                log.info("百炼API流式请求处理完成");
                
            } catch (Exception e) {
                log.error("百炼API请求失败: {}", e.getMessage(), e);
                throw e;
            }
            
            return null; // 返回null表示内容已经通过emitter发送
            
        } catch (Exception e) {
            log.error("调用百炼API时出错", e);
            return "错误：" + e.getMessage();
        }
    }
    
    public String chatCompletion(List<ChatMessage> messages) {
        try {
            // 将消息转换为百炼格式
            List<Map<String, Object>> bailianMessages = convertToBailianFormat(messages);
            
            // 创建请求体
            Map<String, Object> requestBody = createRequestBody(bailianMessages, false);
            
            // 设置请求头
            HttpHeaders headers = createHeaders();
            HttpEntity<Map<String, Object>> requestEntity = new HttpEntity<>(requestBody, headers);
            
            // 百炼API URL
            String url = "https://" + aiConfig.getEndpoint() + "/api/v1/services/aigc/text-generation/generation";
            log.info("发送非流式请求到百炼API: {}", url);
            
            // 发送请求
            Map<String, Object> response = restTemplate.postForObject(url, requestEntity, Map.class);
            
            // 解析百炼响应
            if (response != null && response.containsKey("output") && ((Map)response.get("output")).containsKey("text")) {
                return ((Map<String, String>)response.get("output")).get("text");
            }
            
            return "无法从百炼AI获取响应";
            
        } catch (Exception e) {
            log.error("调用百炼API时出错", e);
            return "错误：" + e.getMessage();
        }
    }
    
    @Override
    public String getModelId() {
        return aiConfig.getModelId();
    }
    
    /**
     * 将内部消息格式转换为百炼API所需格式
     */
    private List<Map<String, Object>> convertToBailianFormat(List<ChatMessage> messages) {
        List<Map<String, Object>> bailianMessages = new ArrayList<>();
        
        for (ChatMessage message : messages) {
            Map<String, Object> bailianMessage = new HashMap<>();
            
            // 百炼API使用system, user, assistant作为角色
            bailianMessage.put("role", message.getRole());
            bailianMessage.put("content", message.getContent());
            
            bailianMessages.add(bailianMessage);
        }
        
        return bailianMessages;
    }
    
    /**
     * 创建请求体
     */
    private Map<String, Object> createRequestBody(List<Map<String, Object>> messages, boolean stream) {
        Map<String, Object> parameters = new HashMap<>();
        parameters.put("max_tokens", aiConfig.getMaxTokens());
        parameters.put("temperature", aiConfig.getTemperature());
        parameters.put("stream", stream);
        
        Map<String, Object> input = new HashMap<>();
        input.put("messages", messages);
        
        Map<String, Object> requestBody = new HashMap<>();
        requestBody.put("model", aiConfig.getModelId());
        requestBody.put("input", input);
        requestBody.put("parameters", parameters);
        
        return requestBody;
    }
    
    /**
     * 创建HTTP请求头
     */
    private HttpHeaders createHeaders() {
        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);
        headers.set("Authorization", "Bearer " + aiConfig.getApiKey());
        headers.set("X-DashScope-SSE", "enable");
        headers.set("X-DashScope-Async", "disable");
        
        return headers;
    }
} 